"grid.py"

from numpy import *
from warnings import warn

from ocemopy.grid import GreatCircle
from ocemopy.tools import PolyGeom

try:
    from matplotlib.toolkits.basemap import Basemap
except:
    warn("Can't import Basemap, so cannot use geographic grids.")

_CGrid_doc = '''
A class for a curvilinear Arikawa C-Grid.
    
    The basis for the CGrid class are two arrays defining the verticies of the
    grid in either Cartesian or geographic coordinates and an optional mask
    defined on the cell centers. Other grid properties, such as the locations
    of the cell centers (rho-points), cell edges (u- and v-points), cell
    widths (dx and dy) and other metrics (angle, dmde, and dndx) are all
    calculated internally from the vertex points. All calculations are done at
    the time the attribute is requested, so changes to the verticies (or mask)
    are instantly recognized by the derived grid properties.
    
    For a geographic grid, a projection may be specified, or The default
    projection for will be defined by the matplotlib.toolkits.Basemap
    projection:
    
        proj = Basemap(projection='merc', resolution=None, lat_ts=0.0)
    
    For a geographic grid, the cell widths are determined by the great circle
    distances.  Angles, however, are defined using the projected coordinates, 
    so a projection that conserves angles must be used.  This means typically
    either Mercator (projection='merc') or Lambert Conformal Conic
    (projection='lcc').
    
    For a grid in Cartesian coordinates:
    
        grd = CGrid(x, y, mask=None, f=None, h=None)

    For a grid in geographic coordinates:
    
        grd = CGrid(lon, lat, proj=<Basemap Mercator projection>, 
                    mask=None, f=None, h=None)
    
    Input vertex arrays may be either numpy arrays or MaskedArrays.  If masked
    arrays are used, the mask will be a combination of the specified mask (if
    given) and the masked locations.
    

'''
class CGrid(object):
    
    __doc__ = _CGrid_doc
    
    def __init__(self, x=None, y=None,  \
                       lon=None, lat=None, \
                       mask=None, proj=None, \
                       f=None, h=None):
        
        # check for neccessary inputs
        assert  (x is not None and y is not None) or \
                (lon is not None and lat is not None),  \
                'Must define either x/y or lat/lon grid verticies'
        
        self.x_vert = x
        self.y_vert = y
        self.lon_vert = lon
        self.lat_vert = lat
        
        # If lat and lon are input, the grid is geographic.
        self.geographic=False
        if  lon is not None and lat is not None:
            self.geographic=True
            if proj is None:    # Use Mercator projection if none given.
                self.proj = Basemap(projection='merc', resolution=None, lat_ts=0.0)
            else:
                self.proj = proj
        
        # Generate the grid mask
        if mask is not None:
            self.mask_rho = mask
        else:
            self.mask_rho = ones(self.x_rho.shape, dtype='d')
        
        # If maskedarray is given for verticies, modify the mask such that 
        # non-existant grid points are masked.  A cell requires all four
        # verticies to be defined as a water point.
        if isinstance(self.x_vert, ma.MaskedArray):
            mask = (self.x_vert.mask[:-1,:-1] | self.x_vert.mask[1:,:-1] | \
                    self.x_vert.mask[:-1,1:] | self.x_vert.mask[1:,1:])
            self.mask_rho = asarray(~(~bool_(self.mask_rho) | mask), dtype='d')
        
        if isinstance(self.y_vert, ma.MaskedArray):
            mask = (self.y_vert.mask[:-1,:-1] | self.y_vert.mask[1:,:-1] | \
                    self.y_vert.mask[:-1,1:] | self.y_vert.mask[1:,1:])
            self.mask_rho = asarray(~(~bool_(self.mask_rho) | mask), dtype='d')
        
        if isinstance(self.lon_vert, ma.MaskedArray):
            mask = (self.lon_vert.mask[:-1,:-1] | self.lon_vert.mask[1:,:-1] | \
                    self.lon_vert.mask[:-1,1:] | self.lon_vert.mask[1:,1:])
            self.mask_rho = asarray(~(~bool_(self.mask_rho) | mask), dtype='d')
        
        if isinstance(self.lat_vert, ma.MaskedArray):
            mask = (self.lat_vert.mask[:-1,:-1] | self.lat_vert.mask[1:,:-1] | \
                    self.lat_vert.mask[:-1,1:] | self.lat_vert.mask[1:,1:])
            self.mask_rho = asarray(~(~bool_(self.mask_rho) | mask), dtype='d')
        
        if (self.x_vert is None or self.y_vert is None) and self.geographic:
            self.calculate_projection()
        
        self.calculate_metrics()
        
        if self.geographic and f is None:
            self.f = 2 * 7.29e-5 * cos(self.lat_rho * pi / 180.)
        else:
            self.f = asarray(f, dtype='d')
            
        self.h = h
    
    def _get_mask_u(self):
        return self.mask_rho[:,1:]*self.mask_rho[:,:-1]
    
    def _get_mask_v(self):
        return self.mask_rho[1:,:]*self.mask_rho[:-1,:]
    
    def _get_mask_psi(self):
        return self.mask_rho[1:,1:]*self.mask_rho[:-1,1:]* \
               self.mask_rho[1:,:-1]*self.mask_rho[:-1,:-1]
    
    def _get_x_rho(self):
        return 0.25*(self.x_vert[1:,1:]+self.x_vert[1:,:-1]+ \
                     self.x_vert[:-1,1:]+self.x_vert[:-1,:-1])
    
    def _get_y_rho(self):
        return 0.25*(self.y_vert[1:,1:]+self.y_vert[1:,:-1]+ \
                     self.y_vert[:-1,1:]+self.y_vert[:-1,:-1])
    
    def _get_x_u(self):
        return 0.5*(self.x_vert[:-1,1:-1] + self.x_vert[1:,1:-1])
    
    def _get_y_u(self):
        return 0.5*(self.y_vert[:-1,1:-1] + self.y_vert[1:,1:-1])
    
    def _get_x_v(self):
        return 0.5*(self.x_vert[1:-1,:-1] + self.x_vert[1:-1,1:])
    
    def _get_y_v(self):
        return 0.5*(self.y_vert[1:-1,:-1] + self.y_vert[1:-1,1:])
    
    def _get_x_psi(self):
        return self.x_vert[1:-1,1:-1]
    
    def _get_y_psi(self):
        return self.y_vert[1:-1,1:-1]
    
    def _get_lon_rho(self):
        if self.lon_vert is None or self.lat_vert is None: return
        return 0.25*(self.lon_vert[1:,1:]+self.lon_vert[1:,:-1]+ \
                     self.lon_vert[:-1,1:]+self.lon_vert[:-1,:-1])
    
    def _get_lat_rho(self):
        if self.lon_vert is None or self.lat_vert is None: return
        return 0.25*(self.lat_vert[1:,1:]+self.lat_vert[1:,:-1]+ \
                     self.lat_vert[:-1,1:]+self.lat_vert[:-1,:-1])
    
    def _get_lon_u(self):
        if self.lon_vert is None or self.lat_vert is None: return
        return 0.5*(self.lon_vert[:-1,1:-1] + self.lon_vert[1:,1:-1])
    
    def _get_lat_u(self):
        if self.lon_vert is None or self.lat_vert is None: return
        return 0.5*(self.lat_vert[:-1,1:-1] + self.lat_vert[1:,1:-1])
    
    def _get_lon_v(self):
        if self.lon_vert is None or self.lat_vert is None: return
        return 0.5*(self.lon_vert[1:-1,:-1] + self.lon_vert[1:-1,1:])
    
    def _get_lat_v(self):
        if self.lon_vert is None or self.lat_vert is None: return
        return 0.5*(self.lat_vert[1:-1,:-1] + self.lat_vert[1:-1,1:])
    
    def _get_lon_psi(self):
        if self.lon_vert is None or self.lat_vert is None: return
        return self.lon_vert[1:-1,1:-1]
    
    def _get_lat_psi(self):
        if self.lon_vert is None or self.lat_vert is None: return
        return self.lat_vert[1:-1,1:-1]
    
    
    def _get_sc_w(self):
        if None in (self.theta_s, self.theta_b, self.hc, self.N): return
        return mgrid[-1.0:0.0:1j*(self.N+1)]
    
    def _get_sc_r(self):
        if None in (self.theta_s, self.theta_b, self.hc, self.N): return
        sc_w = mgrid[-1.0:0.0:1j*(self.N+1)]
        return 0.5*(sc_w[1:]+sc_w[:-1])
    
    def _get_Cs_r(self):
        if None in (self.theta_s, self.theta_b, self.hc, self.N): return
        if self.theta_s == 0.0: return self._get_sc_r()
        return (1-self.theta_b)*sinh(self.theta_s*self._get_sc_r())/ \
               sinh(self.theta_s)+0.5*self.theta_b \
               *(tanh(self.theta_s*(self._get_sc_r()+0.5)) \
               - tanh(0.5*self.theta_s))/tanh(0.5*self.theta_s)
    
    def _get_Cs_w(self):
        if None in (self.theta_s, self.theta_b, self.hc, self.N): return
        if self.theta_s == 0.0: return self._get_sc_w()
        return (1-self.theta_b)*sinh(self.theta_s*self._get_sc_w())/ \
               sinh(self.theta_s)+0.5*self.theta_b \
               *(tanh(self.theta_s*(self._get_sc_w()+0.5)) \
                - tanh(0.5*self.theta_s))/tanh(0.5*self.theta_s)
    
    
    def calculate_metrics(self):
        'Calculates pm, pn, dndx, dmde, and angle from x_vert and y_vert'
        if self.geographic:
            gc_dist = vectorize(lambda lon1, lat1, lon2, lat2: \
                      pyroms.GreatCircle(6378137.0, 6356752.3142, \
                                          lon1, lat1, lon2, lat2).distance)
            lon_temp = 0.5*(self.lon_vert[1:,:]+self.lon_vert[:-1,:])
            lat_temp = 0.5*(self.lat_vert[1:,:]+self.lat_vert[:-1,:])
            if isinstance(lat_temp, ma.MaskedArray): lat_temp = lat_temp.filled(0.0)
            if isinstance(lon_temp, ma.MaskedArray): lon_temp = lon_temp.filled(0.0)
            self.pm = 1.0 / gc_dist(lon_temp[:,1:],  lat_temp[:,1:], \
                                    lon_temp[:,:-1], lat_temp[:,:-1])
            lon_temp = 0.5*(self.lon_vert[:,1:]+self.lon_vert[:,:-1])
            lat_temp = 0.5*(self.lat_vert[:,1:]+self.lat_vert[:,:-1])
            if isinstance(lat_temp, ma.MaskedArray): lat_temp = lat_temp.filled(0.0)
            if isinstance(lon_temp, ma.MaskedArray): lon_temp = lon_temp.filled(0.0)
            self.pn = 1.0 / gc_dist(lon_temp[1:,:],  lat_temp[1:,:], \
                                    lon_temp[:-1,:], lat_temp[:-1,:])
        else:
            x_temp = 0.5*(self.x_vert[1:,:]+self.x_vert[:-1,:])
            y_temp = 0.5*(self.y_vert[1:,:]+self.y_vert[:-1,:])
            self.pm = 1.0 / sqrt(diff(x_temp, axis=1)**2 + diff(y_temp, axis=1)**2)
            x_temp = 0.5*(self.x_vert[:,1:]+self.x_vert[:,:-1])
            y_temp = 0.5*(self.y_vert[:,1:]+self.y_vert[:,:-1])
            self.pn = 1.0 / sqrt(diff(x_temp, axis=0)**2 + diff(y_temp, axis=0)**2)
        
        if any(~isfinite(self.pm)) or any(~isfinite(self.pm)):
            self.pm = ma.masked_where(~isfinite(self.pm), self.pm)
            self.pn = ma.masked_where(~isfinite(self.pn), self.pn)
        
        
        
        self.pm = extrapolate_mask(self.pm, mask=(self.mask==0.0))
        self.pn = extrapolate_mask(self.pn, mask=(self.mask==0.0))
        
        if isinstance(self.pn, ma.MaskedArray):
            self.dndx = ma.zeros((self.Mp, self.Lp), dtype='d')
        else:
            self.dndx = zeros((self.Mp, self.Lp), dtype='d')
        
        if isinstance(self.pm, ma.MaskedArray):
            self.dmde = ma.zeros((self.Mp, self.Lp), dtype='d')
        else:
            self.dmde = zeros((self.Mp, self.Lp), dtype='d')
        
        self.dndx[1:-1,1:-1] = 0.5*(1.0/self.pn[1:-1,2:] - 1.0/self.pn[1:-1,:-2])
        self.dmde[1:-1,1:-1] = 0.5*(1.0/self.pm[2:,1:-1] - 1.0/self.pm[:-2,1:-1])
        
        if self.x_vert is None or self.y_vert is None:
            self.calculate_projection()
        
        self.angle = arctan2(diff(0.5*(self.y_vert[1:,:]+self.y_vert[:-1,:])), \
                           diff(0.5*(self.x_vert[1:,:]+self.x_vert[:-1,:])))
        
        self.angle = extrapolate_mask(self.angle, mask=(self.mask==0.0))
    
    def calculate_projection(self, proj=None):        
        if isinstance(self.lat_vert, ma.MaskedArray):
            mask_lat = self.lat_vert.mask 
            lat_temp = self.lat_vert.filled(0.0)
        else:
            lat_temp = self.lat_vert
        
        if isinstance(self.lon_vert, ma.MaskedArray): 
            mask_lon = self.lon_vert.mask
            lon_temp = self.lon_vert.filled(0.0)
        else:
            lon_temp = self.lon_vert
        
        self.x_vert, self.y_vert = self.proj(lon_temp, lat_temp)
        
        if isinstance(self.lon_vert, ma.MaskedArray):
            self.x_vert = ma.masked_array(self.x_vert, mask=mask_lon)
        
        if isinstance(self.lat_vert, ma.MaskedArray):
            self.y_vert = ma.masked_array(self.y_vert, mask=mask_lat)
    
    def calculate_orthogonality(self):
        '''
        Calculate orthogonality error in radiens
        '''
        z = self.x_vert + 1j*self.y_vert
        du = diff(z, axis=1); du = (du/abs(du))[:-1,:]
        dv = diff(z, axis=0); dv = (dv/abs(dv))[:,:-1]
        ang1 = arccos(du.real*dv.real + du.imag*dv.imag)
        du = diff(z, axis=1); du = (du/abs(du))[1:,:]
        dv = diff(z, axis=0); dv = (dv/abs(dv))[:,:-1]
        ang2 = arccos(du.real*dv.real + du.imag*dv.imag)
        du = diff(z, axis=1); du = (du/abs(du))[:-1,:]
        dv = diff(z, axis=0); dv = (dv/abs(dv))[:,1:]
        ang3 = arccos(du.real*dv.real + du.imag*dv.imag)
        du = diff(z, axis=1); du = (du/abs(du))[1:,:]
        dv = diff(z, axis=0); dv = (dv/abs(dv))[:,1:]
        ang4 = arccos(du.real*dv.real + du.imag*dv.imag)
        ang = mean([abs(ang1), abs(ang2), abs(ang3), abs(ang4)], axis=0)
        ang = (ang-pi/2.0)
        return ang
    
    def maskpoly(self, polyverts, inverse=False, geographic=None):
        """
        Mask Cartesian points contained within the polygons contained in the
        list'polygons'.
        
        A cell is masked if the cell center (x_rho, y_rho) is within the
        polygon. Other sub-masks (mask_u, mask_v, and mask_psi) are updated
        automatically.
        """
        if geographic is None:
            geographic = self.geographic
        
        mask = self.mask_rho
        
        if inverse:
            mask = asarray(~bool_(mask), dtype='d')
        
        iwater = mask == 1.0
        if geographic:
            x_wet = self._get_lon_rho()[iwater]
            y_wet = self._get_lat_rho()[iwater]
        else:
            x_wet = self._get_x_rho()[iwater]
            y_wet = self._get_y_rho()[iwater]
        
        mask_wet = mask[iwater]
        
        inside = PolyGeom(polyverts).inside(zip(x_wet, y_wet))
        
        if any(inside):
            mask_wet[inside] = 0.0
            mask[iwater] = mask_wet
            if inverse:
                mask = asarray(~bool_(a), dtype='d')
            self.mask_rho = mask
    
    
    x = property(lambda self: self.x_vert)
    y = property(lambda self: self.y_vert)
    lon = property(lambda self: self.lon_vert)
    lat = property(lambda self: self.lat_vert)
    mask = property(lambda self: self.mask_rho)
    
    mask_u   = property(_get_mask_u)
    mask_v   = property(_get_mask_v)
    mask_psi = property(_get_mask_psi)
    x_rho = property(_get_x_rho)
    x_u   = property(_get_x_u)
    x_v   = property(_get_x_v)
    x_psi = property(_get_x_psi)
    y_rho = property(_get_y_rho)
    y_u   = property(_get_y_u)
    y_v   = property(_get_y_v)
    y_psi = property(_get_y_psi)
    lon_rho = property(_get_lon_rho)
    lon_u   = property(_get_lon_u)
    lon_v   = property(_get_lon_v)
    lon_psi = property(_get_lon_psi)
    lat_rho = property(_get_lat_rho)
    lat_u   = property(_get_lat_u)
    lat_v   = property(_get_lat_v)
    lat_psi = property(_get_lat_psi)
    

if __name__ == '__main__':
    pass
    # TODO: Write a simple demonstration here...

